-- | Trees

module Tree (Tree.T(..), width, fold, rev, countLfs, dim, depth,
  sparsity, totalDegree, jagged, genNonJagged) where

import Interlude
import Data.Foldable (Foldable)
import qualified Test.QuickCheck as QC

data T a = Lf a | Br [T a] deriving (Eq, Foldable)

instance Show a => Show (Tree.T a) where
  show (Lf a) = show a
  show (Br b) = "{" ++ concat (intersperse ", " $ map show b) ++ "}"

instance Functor Tree.T where
  fmap f (Lf x) = Lf $! f x
  fmap f (Br x) = Br $! map' (fmap f) x

instance {-Add-}Num a => {-Add-}Num (Tree.T a) where
  {-# SPECIALIZE instance Num (Tree.T Double) #-}
  Lf x       + Lf y       = Lf $! x+y
  Br x       + Br y       = Br $! zipKeep (+) x y
  x          + Br (y0:y') = Br (x+y0 !: y')
  Br (x0:x') + y          = Br (x0+y !: x')

  negate = fmap negate
  (*) = error "(*)[Tree.T]"
  abs = error "abs[Tree.T]"
  signum = error "signum[Tree.T]"
  fromInteger = error "fromInteger[Tree.T]"

fold :: (a -> a -> a) -> Tree.T a -> a
fold _ (Lf c) = c
fold f (Br b) = foldl1' f (map' (fold f) b)

rev :: Tree.T a -> Tree.T a
rev (Br b) = Br $! reverse b
rev x = x

width :: Tree.T t -> Int
width (Lf _) = 1
width (Br b) = length b

depth :: T a -> Int
depth (Lf _) = 0
depth (Br x) = 1 + maximum (0 : map depth x)

dim :: Tree.T a -> [Int]
dim (Lf _) = []
dim (Br x) = length x : (map maximum . transpose . map dim) x

sparsity :: Tree.T a -> Double
sparsity x = fromIntegral (countLfs x) / fromIntegral (product $ dim x)

countLfs :: Tree.T t -> Int
countLfs (Lf _) = 1
countLfs (Br b) = sum $ map countLfs b

totalDegree :: Tree.T a -> Int
totalDegree (Lf _) = 0
totalDegree (Br b ) = maximum $ zipWith (\p d -> totalDegree p + d) b [0..]

jagged :: Tree.T a -> Bool
jagged (Lf _) = False
jagged (Br x) = length (nub (map len x)) > 1 || any jagged x where
  len (Lf _) = -1
  len (Br b) = length b

instance (Num a, Arbitrary a) => Arbitrary (T a) where
  arbitrary = gen arbitrary
  shrink = Tree.shrink

shrink :: Num t => T t -> [T t]
shrink (Lf _)  = []
shrink (Br []) = [Lf 0, Lf 1]
shrink (Br x)  = x ++ concat (map Tree.shrink x)

-- | Generator of arbitrary (possibly jagged) trees
gen :: QC.Gen a -> QC.Gen (T a)
gen g = QC.sized jg where
  jg 0 = Lf <$> g
  jg 1 = Lf <$> g
  jg n = randomPart n >>= \p -> Br <$> sequence (map jg p)

-- | Generator of non-jagged trees. (This and the other generators in
--   this module do not have a uniform probability distribution)
genNonJagged :: QC.Gen a -> QC.Gen (T a)
genNonJagged g =
    QC.sized $ \n -> QC.choose (1, (round . (log :: Double -> Double) . fromIntegral) n) >>=
    \k -> QC.vectorOf k (QC.choose (1, round (fromIntegral n ** (1/fromIntegral k) :: Double))) >>= nj where
  nj []     = Lf <$> g
  nj (d:ds) = Br <$> QC.vectorOf d (nj ds)

-- | Generator of random partitions
randomPart :: Int -> QC.Gen [Int]
randomPart 0 = return []
randomPart n = QC.choose (1,n) >>= \k -> (k:) <$> randomPart (n-k)
